import pathlib

import torch.utils.data

import torch_plus
from simple_pendulum_problem import SimplePendulumProblem


class SimplePendulumDataset:
    def __init__(self,
                 problem: SimplePendulumProblem,
                 t_batch: torch.Tensor):
        self._problem = problem
        self._t = torch_plus.clone_detached(t_batch)

        self._t_min = torch.min(self._t)
        self._t_max = torch.max(self._t)

        theta_batch = [problem.theta(float(t)) for t in self._t]
        theta_batch = torch.tensor(theta_batch)
        self._theta = theta_batch

    def problem(self) -> SimplePendulumProblem:
        return self._problem

    def t(self):
        return self._t

    def t_min(self):
        return self._t_min

    def t_max(self):
        return self._t_max

    def theta(self):
        return self._theta

    def save(self, path: pathlib.Path):
        torch.save({
            "t": self._t,
            "problem": {
                "g": self._problem.g(),
                "gamma": self._problem.gamma(),
                "theta0": self._problem.theta0(),
                "m": self._problem.m(),
                "l": self._problem.l_(),
                "t_step": self._problem.t_step()
            }
        }, path)

    @staticmethod
    def load(path: pathlib.Path) -> 'SimplePendulumDataset':
        data = torch.load(path)
        problem = data["problem"]
        return SimplePendulumDataset(
            SimplePendulumProblem(
                problem["g"],
                problem["gamma"],
                problem["theta0"],
                problem["m"],
                problem["l"],
                problem["t_step"]
            ),
            data["t"]
        )
